class Trie:
    class _Node:
        __slots__ = "value", "weight", "children"

        def __init__(self):
            self.value = None
            self.weight = 0
            self.children = {}

        def __contains__(self, ch):
            return ch in self.children

        def __getitem__(self, ch):
            return self.children[ch]

        def __setitem__(self, ch, value):
            self.children[ch] = value

    def __init__(self):
        self.root = self._Node()

    def add(self, word):
        res = 0
        node = self.root
        for ch in word:
            if ch not in node:
                node[ch] = self._Node()
                res += 1
            node = node[ch]
        return res


class Solution:
    def countDistinct(self, s: str) -> int:
        ans = 0
        trie = Trie()
        for i in range(len(s)):
            ans += trie.add(s[i:])
        return ans


if __name__ == "__main__":
    print(Solution().countDistinct("aabbaba"))  # 21
    print(Solution().countDistinct("abcdefg"))  # 28
